#! /usr/bin/python3
# QRNG Appliance evaluation kit command line tool
#
# Copyright (c) Crypta Labs 2022
# Originated by davide@cryptalabs.com 20/07/2022
import socket
import argparse
import struct

verbose = False

def debug_print(str):
  if(verbose):
    print(str)

MAX_BLOCK_SIZE = 4096

class comm_tcp :
  """TCP communication interface"""
  TCP_PORT = 54936
  FLUSH_SIZE = 8000
  TCP_MAX_SIZE = 1072

  def __init__(self, device_name):
    debug_print(socket.gethostbyname(device_name))

    self._sock = socket.socket(socket.AF_INET, # Internet
                     socket.SOCK_STREAM) #TCP
    self._device_name = socket.gethostbyname(device_name)
    self._reconnect()

    return

  def _reconnect(self):
    self._sock.connect((self._device_name, self.TCP_PORT))
    self.set_timeout(1)
    self._rx_buffer = bytearray()
    return

  def flush(self):
  #  self._sock.settimeout(0.000001)
  #  while True:
  #    try:
  #      data, addr = self._sock.recvfrom(self.FLUSH_SIZE)
  #      debug_print("flush:{}".format(len(data)))
  #    except socket.timeout:
  #      # We expect a timeout
  #      break
  #  self._sock.settimeout(self._timeout)
  #  self._rx_buffer = bytearray()
    pass

  def set_timeout(self, timeout):
    self._timeout = timeout
    self._sock.settimeout(self._timeout)
    pass

  def write(self, data):
    return self._sock.sendall(data)

  def read(self, length):
    debug_print("Read")
    while(len(self._rx_buffer) < length):
      try:
        if (length - len(self._rx_buffer) < self.TCP_MAX_SIZE):
          to_read = length - len(self._rx_buffer)
        else:
          to_read = self.TCP_MAX_SIZE;
        data = self._sock.recv(to_read)
        if(len(data) == 0):
          ret_data = bytes(self._rx_buffer)
          self._rx_buffer = bytearray()
          return ret_data
        # TODO: check address
        self._rx_buffer.extend(data)
        #print(data)
      except socket.timeout:
        # A timeout can happen:
        #  - Empty Rx buffer
        #  - Return current data
        debug_print("timeout, rx buffer len:{}".format(len(self._rx_buffer)))
        ret_data = bytes(self._rx_buffer)
        self._rx_buffer = bytearray()
        return ret_data
    debug_print("rx buffer len: {}. requested length:{}".format(len(self._rx_buffer),length))
    ret_data = bytes(self._rx_buffer[:length])
    del self._rx_buffer[:length]

    return ret_data

  def close(self):
    self._sock.close()
    pass

class qrng_cmdctrl:

  CMDCTRL_CMD_GET_STATUS = b'\x01'
  CMDCTRL_CMD_START = b'\x04'
  CMDCTRL_CMD_GET_CONFIG = b'\x07'
  CMDCTRL_CMD_SET_CONFIG = b'\x08'
  CMDCTRL_CMD_GET_STATISTICS = b'\x09'
  CMDCTRL_CMD_RESET = b'\x0A'
  CMDCTRL_CMD_INFO = b'\x0B'

  CMDCTRL_RESP_ACK = b'\x11'
  CMDCTRL_RESP_NACK = b'\x12'
  CMDCTRL_RESP_CONFIG = b'\x17'
  CMDCTRL_RESP_STATISTICS = b'\x19'
  CMDCTRL_RESP_INFO = b'\x1B'

  _success_response = {
    CMDCTRL_CMD_GET_STATUS : CMDCTRL_RESP_ACK,
    CMDCTRL_CMD_START : CMDCTRL_RESP_ACK,
    CMDCTRL_CMD_GET_CONFIG : CMDCTRL_RESP_CONFIG,
    CMDCTRL_CMD_SET_CONFIG : CMDCTRL_RESP_ACK,
    CMDCTRL_CMD_GET_STATISTICS : CMDCTRL_RESP_STATISTICS,
    CMDCTRL_CMD_RESET : CMDCTRL_RESP_ACK,
    CMDCTRL_CMD_INFO : CMDCTRL_RESP_INFO,
  }

  _payload_size = {
    CMDCTRL_RESP_ACK : 5,
    CMDCTRL_RESP_NACK : 0,
    CMDCTRL_RESP_CONFIG : 12,
    CMDCTRL_RESP_STATISTICS : 30,
    CMDCTRL_RESP_INFO : 56,
  }

  def __init__(self, device_name):
    self._comm = comm_tcp(device_name)
    self.resp_payload = b'0'

  def start_one_shot(self, length):
    payload = struct.pack("<BH", 0x01, (int)(length))
    if(self.command(self.CMDCTRL_CMD_START, payload)):
      debug_print("")
      self._comm.set_timeout(0.1 + 0.0001*length)
      qrnd = self._comm.read(length)
      if(len(qrnd) == length):
        return qrnd
      else:
        return None
    else:
     return None
  
  def reset(self):
    return self.command(self.CMDCTRL_CMD_RESET)

  def get_status(self):
    if(self.command(self.CMDCTRL_CMD_GET_STATUS, None)):
      fields = struct.unpack("<BI", self.resp_payload)
      return {'started' : (fields[0]) & 1, 
              'startup_test_in_progress' : (fields[0]>>1) & 1,
              'voltage_low' : (fields[0]>>2) & 1, 
              'voltage_high' : (fields[0]>>3) & 1, #
              'voltage_undefined' : (fields[0]>>4) & 1,
              'bitcount' : (fields[0]>>5) & 1, 
              'repetition_count' : (fields[0]>>6) & 1, 
              'adaptive_proportion' : (fields[0]>>7) & 1, 
              'ready_bytes' : fields[1] }
    else:
     return None

  def get_config(self):
    if(self.command(self.CMDCTRL_CMD_GET_CONFIG, None)):
      fields = struct.unpack("<BfBBBHH", self.resp_payload)
      return {'postprocess' : fields[0],
              'inital_level' : fields[1], 
              'startup_test' : (fields[2]) & 1,
              'auto_calibration' : (fields[2] >> 1) & 1, 
              'repetition_count' : (fields[2] >> 2) & 1,
              'adaptive_proportion' : (fields[2] >> 3) & 1,
              'bit_count' : (fields[2] >> 4) & 1,
              'generate_on_error' : (fields[2] >> 5) & 1,
              'n_lsbits' : fields[3], 
              'hash_input_size' : fields[4], 
              'block_size' : fields[5],
              'autocalibration_target' : fields[6]}
    else:
      return None

  def get_statistics(self):
    if(self.command(self.CMDCTRL_CMD_GET_STATISTICS, None)):
      fields = struct.unpack("<QIIIIHf", self.resp_payload)
      return {'generated_bytes' : fields[0],
              'repetition_count_failures' : fields[1], 
              'adaptive_proportion_failures' : fields[2], 
              'bitcount_failures' : fields[3],
              'speed' : fields[4], 
              'sensif_average' : fields[5], 
              'ledctrl_level' : fields[6]}
    else:
      return None

  def set_config_internal(self,
                  postprocess,
                  inital_level,
                  startup_test,
                  auto_calibration,
                  repetition_count,
                  adaptive_proportion,
                  bit_count,
                  generate_on_error,
                  n_lsbits,
                  hash_input_size,
                  block_size,
                  autocalibration_target):

    byte_flags = (startup_test) | \
                  (auto_calibration << 1) | \
                  (repetition_count << 2) | \
                  (adaptive_proportion << 3) | \
                  (bit_count << 4) | \
                  (generate_on_error << 5)

    payload = struct.pack("<BfBBBHH",postprocess, inital_level, byte_flags, n_lsbits, hash_input_size, block_size, autocalibration_target)
    return self.command(self.CMDCTRL_CMD_SET_CONFIG, payload)

  def set_config(self, config):
    return self.set_config_internal( 
          config['postprocess'],
          config['inital_level'],
          config['startup_test'],
          config['auto_calibration'],
          config['repetition_count'],
          config['adaptive_proportion'],
          config['bit_count'],
          config['generate_on_error'],
          config['n_lsbits'],
          config['hash_input_size'],
          config['block_size'],
          config['autocalibration_target'])

  def get_info(self):
    if(self.command(self.CMDCTRL_CMD_INFO, None)):
      fields = struct.unpack("<II24s24s", self.resp_payload)
      return {'core_version' : fields[0], 
              'sw_version' : fields[1],
              'serial' : fields[2][:fields[2].index(0)], 
              'hw_info' : fields[3][:fields[3].index(0)]}
    else:
     return None

  def command(self, cmd_code, payload=None):
    self._comm.flush()
    if(cmd_code not in self._success_response.keys()):
      raise ValueError("Command code 0x{02X} not valid".format(cmd_code))

    to_send = cmd_code; 
    if(payload != None):
      to_send += payload
    
    try:
      self._comm.write(to_send)
    except IOError as ex:
      print("Exception:" + str(ex))
      return False

    try:
      self._comm.set_timeout(1)
      resp = self._comm.read(1)
      debug_print("response length: {}".format(len(resp)))
      if(len(resp) == 0):
        return False
    except Exception as ex:
      debug_print("Exception:" + str(ex))
      return False

    if(resp == self._success_response[cmd_code]):
      expected_payload = self._payload_size[self._success_response[cmd_code]]
      debug_print(expected_payload)
      if(expected_payload > 0):
        try:
          self._comm.set_timeout(0.001*expected_payload)
          self.resp_payload = self._comm.read(expected_payload)
          debug_print("Payload length: {}".format(len(self.resp_payload)))
          if(len(self.resp_payload) != expected_payload):
            return False
        except IOError as ex:
          debug_print("Exception:" + str(ex))
          return False

      return True
    elif (resp == self.CMDCTRL_RESP_NACK):
      return False
    else:
      raise RuntimeError("Unexpected byte received: 0x{}".format(resp.hex()))
    
  def get_resp_payload(self):
    return self.resp_payload

  def close(self):
    return self._comm.close()

# Command line tool
if __name__ == "__main__":

  _raw_config = {
    'postprocess' : 1,
    'inital_level' : 60, 
    'startup_test' : 1,
    'auto_calibration' : 1, 
    'repetition_count' : 1,
    'adaptive_proportion' : 1,
    'bit_count' : 1,
    'generate_on_error' : 0,
    'n_lsbits' : 5, 
    'hash_input_size' : 0, 
    'block_size' : 3360,
    'autocalibration_target' : 320
  }

  _sha256_config = {
    'postprocess' : 0,
    'inital_level' : 60, 
    'startup_test' : 1,
    'auto_calibration' : 1, 
    'repetition_count' : 1,
    'adaptive_proportion' : 1,
    'bit_count' : 1,
    'generate_on_error' : 0,
    'n_lsbits' : 5, 
    'hash_input_size' : 42, 
    'block_size' : 3360,
    'autocalibration_target' : 320
  }

  def postprocess_name_to_code(name):
    if (name == 'RAW_NOISE'):
      return 1
    elif (name == 'SHA256'):
      return 0
    else:
      return 0xff
      
  parser = argparse.ArgumentParser(description= 'QRNG Appliance Command and Control protocol host utility by Crypta Labs')
  parser.add_argument('-d','--device', dest='device',
                      help='Device interface name or IP')
  parser.add_argument('-s','--get-status', dest='status', action = "store_true",
                      help='Get current QRNG configuration')
  parser.add_argument('-g','--get-config', dest='get_config', action = "store_true",
                      help='Get current QRNG conditioning')
  parser.add_argument('-t','--get-statistics', dest='get_stats', action = "store_true",
                      help='Get QRNG statistics')
  parser.add_argument('-r','--read', dest='read_size', type=int,
                      help='One shot read')
  parser.add_argument('-o','--output-file', dest='output_file', type=str,
                      help='Output file (used with -r, --read)')
  parser.add_argument('-a','--append', dest='append', action = "store_true",
                      help='Append to output file (used with -o, --output-file)')
  parser.add_argument('--reset', action = "store_true",
                      help='Reset error statistics and repeat start-up test')
  parser.add_argument('--info', action = "store_true",
                      help='Get device information')

  # QRNG configuration parameters
  parser.add_argument('--postprocess', dest='postprocess', choices=['RAW_NOISE','SHA256'],
                      help='Type of conditioning')
                      
  parser.add_argument('-v','--verbose', action = "store_true",
                      help='Print debug information and random data')

  args = parser.parse_args()

  cc = qrng_cmdctrl(args.device) 

  if(args.verbose):
    verbose = True

  if(args.status):
    status = cc.get_status()
    if(status is not None):
      print("### STATUS:")
      print(" started: {}".format(status['started']) )
      print(" startup_test_in_progress: {}".format(status['startup_test_in_progress']) )
      print(" voltage_low: {}".format(status['voltage_low']) )
      print(" voltage_high: {}".format(status['voltage_high']) )
      print(" voltage_undefined: {}".format(status['voltage_undefined']) )
      print(" bitcount: {}".format(status['bitcount']) )
      print(" repetition_count: {}".format(status['repetition_count']) )
      print(" adaptive_proportion: {}".format(status['adaptive_proportion']) )
      print(" Ready bytes: {}".format(status['ready_bytes']) )
    else:
      print("Error reading status")
      quit()

  if(args.get_config):
    config = cc.get_config()

    if(config is not None):
      post_proc_str ="Unknown"

      if(config['postprocess'] == 0):
        post_proc_str ="SHA256"
      elif(config['postprocess'] == 1):
        post_proc_str ="RAW NOISE"

      print("### Post process Configuration:{}".format(post_proc_str))
      debug_print(" initial_level: {}".format(config['inital_level']) )
      debug_print(" startup_test: {}".format(config['startup_test']) )
      debug_print(" auto_calibration: {}".format(config['auto_calibration']) )
      debug_print(" repetition_count: {}".format(config['repetition_count']) )
      debug_print(" adaptive_proportion: {}".format(config['adaptive_proportion']) )
      debug_print(" bit_count: {}".format(config['bit_count']) )
      debug_print(" generate_on_error: {}".format(config['generate_on_error']) )
      debug_print(" n_lsbits: {}".format(config['n_lsbits']) )
      debug_print(" hash_input_size: {}".format(config['hash_input_size']) )
      debug_print(" block_size: {}".format(config['block_size']) )
      debug_print(" autocalibration_target: {}".format(config['autocalibration_target']) )
    else:
      print("Error reading configuration")
      quit()

  if(args.get_stats):
    stats = cc.get_statistics()
    if(stats is not None):
      print("### STATISTICS:")
      print(" generated bytes: {}".format(stats['generated_bytes']) )
      print(" repetition count failures: {}".format(stats['repetition_count_failures']) )
      print(" adaptive proportion failures: {}".format(stats['adaptive_proportion_failures']) )
      print(" bitcount failures: {}".format(stats['bitcount_failures']) )
      print(" LED level: {}%".format(stats['ledctrl_level']) )
      print(" ADC average: {}".format(stats['sensif_average']) )
    else:
      print("Error reading statistics")
      quit()

  if(args.postprocess):
    if(args.postprocess == 'SHA256'):
      ret = cc.set_config(_sha256_config)
    elif(args.postprocess == 'RAW_NOISE'):
      ret = cc.set_config(_raw_config)
    else:
      print("Unknown postprocess")
      quit()
      
    if(ret):
      print("### SET CONFIGURATION DONE")
    else:
      print("Error during QRNG configuration")
      quit()

  if(args.info):
    info = cc.get_info()
    if(info is not None):
      print("### GET INFO DONE")
      print(" Core version: {}.{}".format(info['core_version'] >> 16, info['core_version'] & 0xffff) )
      print(" SW version: {}.{}".format(info['sw_version'] >> 16, info['sw_version'] & 0xffff) )
      print(" Serial: {}".format(info['serial'].decode("utf-8")))
      print(" HW info: {}".format(info['hw_info'].decode("utf-8")))
    else:
      print("Error reading info")
      quit()

  if(args.reset):
    if(cc.reset() is not None):
      print("Reset done!")
    else:
      print("Reset error!")
      quit()

  if(args.read_size):
    reading = cc.start_one_shot(args.read_size)
    if(reading != None):
      print("One shot read success")
      if(args.output_file):
        with open(args.output_file, "ab" if (args.append) else "wb") as binary_file:
          # Write bytes to file
          binary_file.write(reading)
      debug_print("data:");
      debug_print(reading.hex())
    else:
      print("Error one shot reading")
      quit()

  quit()


